#include "kheap.h"
#include "mm.h"
#include "stdio.h"

static block_head* heap_start;
static u32 heap_total_size;

void print_heap() {
  block_head* cur = heap_start;

  while(cur != END_OF_HEAP) {
    if(cur->magic != HEAP_MAGIC) {
      kprintf("corrupted block at %p\n", (u32)cur);
    }

    else {
      kprintf("block at: %p size: %d", (u32)cur + sizeof(block_head), cur->size);
      
      if(cur->free)
	kprintf(" free\n");
      else
	kprintf(" used\n");
    }
    cur = cur->next;
  }
  return;
}

void init_heap(u32 start, u32 sz) {
  heap_start = (block_head*)start;
  heap_start->magic = HEAP_MAGIC;
  heap_start->size = sz - sizeof(block_head) - sizeof(block_head);
  heap_start->free = TRUE;
  heap_start->prev = 0;
  heap_start->next = (block_head*)(start + sz - sizeof(block_head));
  (heap_start->next)->size = 0;
  (heap_start->next)->free = FALSE;
  (heap_start->next)->magic = HEAP_MAGIC;
  heap_start->next->next = END_OF_HEAP;
  heap_total_size = sz;
}

u32* kmalloc(u32 sz) {
  block_head* cur = heap_start;
  block_head* next;
  block_head* prev = heap_start->prev;
  u32 tmp, dif;  

  // go through all the blocks to see if we can find a free one
  // that is the right size. if next is 0, we are at end of heap
  while(cur->next != END_OF_HEAP) {
    
    // if we find one, expand the next block backwards 
    // to take up the extra slack
    if(cur->size >= (sz + sizeof(block_head)) && cur->free == TRUE) {
      next = cur->next;
      cur->prev = prev;
      cur->next = (block_head*)((u32)cur + sizeof(block_head) + sz);
      cur->next->size = next->size + cur->size - sz;
      cur->next->free = TRUE;
      cur->next->next = next;
      cur->next->prev = cur;
      cur->next->magic = HEAP_MAGIC;
      cur->free = FALSE;
      cur->size = sz;
      return (u32*)((u32)cur + sizeof(block_head));
    }

    prev = cur;
    cur = cur->next;
  }

  return 0;
  // if we got to this point we need to expand the heap
  // heap is expanded a page at a time and never shrinks
}

u32* kmalloc_a(u32 sz) {
  block_head* cur = heap_start;
  block_head* next;
  block_head* prev = heap_start->prev;
  u32 tmp, orig_size;  

  while(cur->next != END_OF_HEAP) {
    if(cur->size >= (sz + sizeof(block_head) + (0x1000 - ((u32)cur & 0x00000fff)))
		     && cur->free == TRUE) {
      tmp = (0x1000 + (u32)cur) & 0xFFFFF000;
      next = cur->next;
      orig_size = cur->size;
      cur->size = (tmp - (u32)cur) - sizeof(block_head);
      cur->next = (block_head*)(tmp - sizeof(block_head));
      cur->next->size = sz;
      cur->next->free = FALSE;
      cur->next->magic = HEAP_MAGIC;
      cur->next->prev = cur;
      cur->next->next = (block_head*)(tmp + sz);
      cur->next->next->magic = HEAP_MAGIC;
      cur->next->next->size = orig_size - sz - cur->size;
      cur->next->next->free = TRUE;
      cur->next->next->prev = cur->next;
      cur->next->next->next = next;
      return (u32*)((u32)cur->next + sizeof(block_head));
    }

    prev = cur;
    cur = cur->next;
  }

  return 0;
}

// frees a block, and combines it with prev and next blocks if they are free
void kfree(u32* block) {
  block_head* cur = (block_head*)((u32)block - sizeof(block_head));
  cur->free = TRUE;

  if(cur->next->free == TRUE) {
    cur->size = cur->size + cur->next->size;// + sizeof(block_head);
    cur->next = cur->next->next;
    cur->next->prev = cur;
  }

  if(cur->prev->free == TRUE && cur->prev != 0) {
    cur->prev->size = cur->prev->size + cur->size;// + sizeof(block_head);
    cur->next->prev = cur->prev;
    cur->prev->next = cur->next;
  }
  
}
